# coding: utf-8
import pickle
import random
from random import shuffle

import numpy as np
import matplotlib.pyplot as plt

from TSP import TSP

np.set_printoptions(suppress=True, linewidth=999, threshold=np.inf)
plt.rc('font', family='Times New Roman', weight='normal', size=12)


def copy(target):
    return pickle.loads(pickle.dumps(target))


class Encoding:
    def __init__(self, data, cost_func):
        self.cities = None
        self.cost = 1e99
        # init params
        self.cities = copy(data.cities)
        self.depot = data.depot
        self.cities.remove(self.depot)
        shuffle(self.cities)
        self.cities.insert(0, self.depot)
        self.cost_func = cost_func
            
    def __repr__(self):
        return f'cities: {self.cities}, total: {self.cost}'
    
    def calc_cost(self):
        self.cost = self.cost_func(route=self.cities + [self.depot])
        return self

    def two_opt(self, distance, depth=1):
        for _ in range(depth):
            for i in range(1, len(self.cities)):
                for j in range(i + 1, len(self.cities)):
                    d1 = distance[self.cities[i - 1]][self.cities[i]] + distance[self.cities[j - 1]][self.cities[j]]
                    d2 = distance[self.cities[i - 1]][self.cities[j - 1]] + distance[self.cities[i]][self.cities[j]]
                    if d2 < d1:
                        self.cities = self.cities[:i] + list(reversed(self.cities[i:j])) + self.cities[j:]
        return self
    
    
# the abstract (x - b) for sequence form
def optimal_indexes(x, b):
    return [b.index(i) for i in x]


# the abstract (v1 - v2) for sequence form
def velocity_merging(v1, v2):
    return np.add(v1, v2) / 2.0


def normalize(v):
    _v = sorted(v)
    return np.array([_v.index(i) for i in v])


def calc_velocity(x, v0, p, g, w, c1, c2):
    r1 = random.uniform(0, 1)
    r2 = random.uniform(0, 1)
    v = velocity_merging(w * v0,
                         velocity_merging(c1 * r1 * np.array(optimal_indexes(x, p)),
                                          np.array(c2 * r2 * np.array(optimal_indexes(x, g))))
                         )
    return normalize(v)


class PSO:
    def __init__(self, data):
        self.data = data
        self.cost = data.cost
        self.NP = 20
        self.w = 1
        self.c1 = 1
        self.c2 = 1
        self.Population = [Encoding(self.data, self.cost)
                           .two_opt(self.data.distance_matrix).calc_cost() for _ in range(self.NP)]
        self.pBest_group = copy(self.Population)
        self.gBest = copy(min(self.Population, key=lambda x: x.cost))
        self.evo_list = []

    def optimize(self, iterations=1000):
        velocity = np.zeros(shape=(self.NP, len(self.Population[0].cities)), dtype=np.int32)
        for it in range(iterations):
            for i in range(self.NP):
                x = self.Population[i].cities
                v = velocity[i]
                pBest = self.pBest_group[i].cities
                v = calc_velocity(x, v, pBest, self.gBest.cities, self.w, self.c1, self.c2)  # calc velocity
                self.Population[i].cities = [x[j] for j in v]  # update position
                self.Population[i].two_opt(distance=self.data.distance_matrix)  # 2-opt
                self.Population[i].cost = self.cost(route=self.Population[i].cities + [self.data.depot])
                # update pBest
                if self.pBest_group[i].cost > self.Population[i].cost:
                    self.pBest_group[i] = copy(self.Population[i])
                # update gBest
                if self.gBest.cost > self.pBest_group[i].cost:
                    self.gBest = copy(self.pBest_group[i])
            
            self.evo_list.append(self.gBest.cost)
            # visualization
            plt.clf()
            plt.subplot(1, 2, 1)
            self.data.draw(plt, routes=[self.gBest.cities])
            plt.title(self.gBest.cost)
            plt.subplot(1, 2, 2)
            plt.plot(range(len(self.evo_list)), self.evo_list, label=f'cost: {self.evo_list[-1]:.2f}')
            plt.title(f'optimal cost')
            plt.pause(0.001)
        plt.show()
        
        
plt.figure(figsize=(10, 4))
tsp = TSP(path=f'../../data/TSPLIB/kroA200.tsp.txt')
optimizer = PSO(data=tsp)
optimizer.optimize()

# x = [1, 2, 3, 4, 5]  # route
# v0 = np.zeros(shape=(len(x)), dtype=np.int32)
# print(v0)
# pBest = [1, 3, 4, 2, 5]
# gBest = [1, 3, 4, 2, 5]
# v = calc_velocity(x, v0, pBest, gBest)
#
# x = [x[i] for i in v]
# print(x)
